import phot


if __name__ == "__main__":
    """双偏振光收发模块 + 光纤信道"""
    """本代码为程序主函数 本代码主要适用于 QPSK，16QAM，64QAM 调制格式的单载波相干背靠背（B2B）信号"""

    phot.config(plot=False, backend="torch")  # 全局开启画图，backend 使用 numpy

    # 设置全局系统仿真参数
    num_symbols = 2**16                              # 符号数目
    bits_per_symbol = 4                              # 2 for QPSK; 4 for 16QAM; 6 for 64QAM  设置调制格式
    total_baud = 60e9                                # 信号波特率，符号率
    up_sampling_factor = 2                           # 上采样倍数
    sampling_rate = up_sampling_factor * total_baud  # 信号采样率
    Reference_frequency = 193.1e12

    # 首先产生发射端X/Y双偏振信号
    signal_bits = phot.gen_bits(num_symbols * bits_per_symbol)  #262144

    # 生成双偏振符号序列 #
    signals = phot.qam_modulate(signal_bits, bits_per_symbol)   #4

    # 均衡的训练序列
    prev_symbols = signals

    # 上采样 #
    signals = phot.up_sample(signals, up_sampling_factor)

    # RRC #
    RRC_ROLL_OFF = 0.01
    signals = phot.pulse_shaper(signals, up_sampling_factor, RRC_ROLL_OFF, total_baud)

    # DAC #
    DAC_Resolution_Bits = 8  # DAC量化bit位数
    DAC_OUTPUT_VPP = 1
    signals = phot.dac(signals, DAC_Resolution_Bits, DAC_OUTPUT_VPP)

    """ 光源 """
    Linewidth = 100e3
    output_power_dB = 0
    laser_center_frequency = Reference_frequency
    Optical_tx_laser = phot.laser_cw(num_symbols*up_sampling_factor, laser_center_frequency, Reference_frequency, sampling_rate, output_power_dB, Linewidth)

    """ 调制器 """
    Extinction_ratio_db = 50
    VPI = 5
    VDC = -2.5
    signals = phot.iq_modulator(signals, VDC, Optical_tx_laser, Extinction_ratio_db, VPI)

    """ 根据OSNR添加高斯白噪声 """
    osnr_db = 20  # 设置系统OSNR，也就是光信号功率与噪声功率的比值，此处单位为dB
    signals = phot.gaussian_noise(signals, osnr_db, sampling_rate)

    """ Optical Fiber Channel, 入纤功率都没有, 光纤建模不够完善 """
    # 实际情况：1000公里 10米一步
    num_spans = 1  # 多少个 span (每个span经过一次放大器)
    span_length = 50  # 一个 span 的长度 (km)
    delta_z = 10  # 单步步长 (km)
    alpha = 0.2
    beta2 = 21.6676e-24
    gamma = 1.3
    signals, signals_power = phot.fiber(signals, sampling_rate, num_spans, beta2, delta_z, gamma, alpha, span_length)

    """" 相干接收机本振激光器 """
    LO_Linewidth = 100e3                                   # 激光器线宽
    frequency_offset = 2e9
    LO_CENTER_Frequency = Reference_frequency + frequency_offset            # LO中心频率
    LO_power_dB = 20                                      # power OF LO
    Optical_LO = phot.laser_cw(num_symbols*up_sampling_factor, LO_CENTER_Frequency, Reference_frequency, sampling_rate, LO_power_dB, LO_Linewidth)

    """" 相干接收机 """
    Responsivity = 1
    signals = phot.coherent_receiver(signals, Optical_LO, Responsivity)              ## 由接收机组件组成的接收机

    """ 模拟接收机造成的I/Q失衡，主要考虑幅度失衡和相位失衡，这里将两者都加在虚部上 """
    signals = phot.add_iq_imbalance(signals)                                ## 接收机的损伤

    """ 加入ADC的量化噪声 """
    adc_sample_rate = 160e9  # ADC采样率
    adc_resolution_bits = 8  # ADC的bit位数
    signals = phot.adc_noise(signals, sampling_rate, adc_sample_rate, adc_resolution_bits)

    """ IQ正交化补偿，就是将之前的I/Q失衡的损伤补偿回来 """
    signals = phot.iq_compensation(signals)   # 自己的方案

    """ 粗糙的频偏估计和补偿 """
    signals = phot.freq_offset_compensation(signals, sampling_rate)

    """ 接收端相应的RRC脉冲整形，具体的参数代码与发射端的RRC滤波器是一致的 """
    signals = phot.pulse_shaper(signals, up_sampling_factor, RRC_ROLL_OFF, total_baud)

    """ 帧同步，寻找与发射端原始信号头部对应的符号 """
    signals, prev_symbols = phot.synchronization(signals, prev_symbols, up_sampling_factor)

    """ 自适应均衡，此处采用恒模算法（CMA）对收敛系数进行预收敛，再拿收敛后的滤波器系数对正式的信号使用半径定向算法（RDE）进行均衡收敛，"""
    num_tap = 25  # 均衡器抽头数目，此处均衡器内部是采用FIR滤波器，具体可查阅百度或者论文，
    ref_power_cma = 4  # 设置CMA算法的模
    cma_convergence = 30000  # CMA预均衡收敛的信号长度
    step_size_cma = 1e-9  # CMA的更新步长，梯度下降法的步长
    step_size_rde = 1e-9  # RDE的更新步长，梯度下降法的步长，%% CMA和RDE主要就是损失函数不同
    signals = phot.adaptive_equalizer(
        signals,
        num_tap,
        cma_convergence,
        ref_power_cma,
        step_size_cma,
        step_size_rde,
        up_sampling_factor,
        bits_per_symbol,
        total_baud,
    )

    """ 均衡后进行精确的频偏估计和补偿 采用FFT-FOE算法，"""
    signals = phot.freq_offset_compensation(signals, total_baud)

    """ 相位恢复  采用盲相位搜索算法（BPS）进行相位估计和补偿 """
    num_test_angle = 64  # BPS算法的测试角数目
    block_size = 100  # BPS算法的块长设置
    signals = phot.bps_restore(signals, num_test_angle, block_size, bits_per_symbol)

    # 分析器画星座图
    phot.constellation_diagram(signals)
    #
    # # 分析器画眼图
    phot.eye_diagram(signals, up_sampling_factor)

    # 此处开始计算误码率
    ber, q_factor = phot.bits_error_count(signals, prev_symbols, bits_per_symbol)
    print(ber)
    print(q_factor)
